package vllm

import (
	"embed"
	"fmt"
	"golang.org/x/net/context"
)

//go:embed chat_template.tmpl
var chatTemplateFS embed.FS

type PredictOptions struct {
	Model        string    `json:"model"`
	Messages     []Message `json:"messages"`
	MaxTokens    int       `json:"max_tokens"`
	Temperature  float64   `json:"temperature"`
	TopP         float64   `json:"top_p"`
	StopWords    []string  `json:"stop_words"`
	Stream       bool      `json:"stream"`
	ChatTemplate string    `json:"chat_template"`
}

type PredictOption func(p *PredictOptions)

var DefaultPredictOptions = PredictOptions{
	MaxTokens:   512,
	Temperature: 0.7,
	TopP:        1.0,
	Stream:      false,
}

func GetChatTemplate(fileName string) (string, error) {
	data, err := chatTemplateFS.ReadFile(fileName)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func SetModel(model string) PredictOption {
	return func(p *PredictOptions) {
		p.Model = model
	}
}

func SetMessages(messages []Message) PredictOption {
	return func(p *PredictOptions) {
		p.Messages = messages
	}
}

func SetMaxTokens(maxTokens int) PredictOption {
	return func(p *PredictOptions) {
		p.MaxTokens = maxTokens
	}
}

func SetTemperature(temperature float64) PredictOption {
	return func(p *PredictOptions) {
		p.Temperature = temperature
	}
}

func SetTopP(topP float64) PredictOption {
	return func(p *PredictOptions) {
		p.TopP = topP
	}
}

func SetStopWords(stop []string) PredictOption {
	return func(p *PredictOptions) {
		p.StopWords = stop
	}
}

func SetStream(stream bool) PredictOption {
	return func(p *PredictOptions) {
		p.Stream = stream
	}
}

func SetChatTemplate(chatTemplate string) PredictOption {
	return func(p *PredictOptions) {
		p.ChatTemplate = chatTemplate
	}
}

func NewPredictOptions(opts ...PredictOption) PredictOptions {
	options := DefaultPredictOptions
	for _, opt := range opts {
		opt(&options)
	}
	if options.ChatTemplate == "" {
		options.ChatTemplate, _ = GetChatTemplate("chat_template.tmpl")
	}
	fmt.Printf("chat_template:%s", options.ChatTemplate)
	return options
}

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type VLLM struct {
	client    *VLLMClient
	modelPath string
}

func NewVLLM(url, model string) *VLLM {
	return &VLLM{
		client:    NewVLLMClient(url, model),
		modelPath: model,
	}
}

func (v *VLLM) PredictByVLLM(opts ...PredictOption) (*Predict, error) {
	options := []PredictOption{
		SetModel(v.client.modelPath),
	}
	options = append(options, opts...)
	return v.client.Predict(context.Background(), options...)
}
